﻿#pragma once
#include "DataObject.h"
#include "IsoVortexIC.h"
#include "MeshBase.h"
#include "SystemBase.h"
#include "Types.h"

#include <vtkm/Matrix.h>
#include <vtkm/TypeTraits.h>
#include <vtkm/cont/Timer.h>
#include <vtkm/io/VTKDataSetWriter.h>
#include <vtkm/worklet/WorkletCellNeighborhood.h>
#include <vtkm/worklet/WorkletMapField.h>
#include <vtkm/worklet/WorkletMapTopology.h>

#include <fmt/color.h>
#include <fmt/format.h>

template<UInt8 NUM_VARIABLES>
struct DataPack
{
  VTKM_EXEC
  DataPack()
  {
    for (size_t i = 0; i < NUM_VARIABLES; i++)
    {
      duh[i] = vtkm::TypeTraits<Vec3f>::ZeroInitialization();
    }
  }
  Vec<FloatDefault, NUM_VARIABLES> uh;
  Vec<Vec3f, NUM_VARIABLES> duh;
  //Vec3f duh[NUM_VARIABLES];
  Vec3f inv_term[NUM_VARIABLES]; // 无粘项
  Vec3f vis_term[NUM_VARIABLES]; // 粘性项

  VTKM_EXEC
  void InvisTerm()
  {
    FloatDefault _gamma = 1.4;
    Vec3f vel{ uh[1] / uh[0], uh[2] / uh[0], uh[3] / uh[0] };

    FloatDefault p =
      (_gamma - 1) * (uh[4] - 0.5 * (uh[1] * uh[1] + uh[2] * uh[2] + uh[3] * uh[3]) / uh[0]);
    FloatDefault h = (uh[4] + p) / uh[0];

    inv_term[0] = uh[0] * vel;

    inv_term[1] = uh[1] * vel;
    inv_term[1][0] += p;

    inv_term[2] = uh[2] * vel;
    inv_term[2][1] += p;

    inv_term[3] = uh[3] * vel;
    inv_term[3][2] += p;

    inv_term[4] = uh[0] * h * vel;
  }

  VTKM_EXEC
  void ViscousTerm()
  {
    FloatDefault _reynolds = 1000, _prandtl = 0.72;
    FloatDefault _gamma = 1.4;

    using RealTensor = vtkm::Matrix<FloatDefault, 3, 3>;
    //using RealTensor = Vec<Vec3f, 3>;
    auto r = uh[0];
    auto drho = duh[0];
    Vec3f vel{ uh[1] / uh[0], uh[2] / uh[0], uh[3] / uh[0] };


    RealTensor grad_mom;
    //{ duh[1], duh[2], duh[3] };
    grad_mom[0] = duh[1];
    grad_mom[1] = duh[2];
    grad_mom[2] = duh[3];

    RealTensor grad_vel;
    grad_vel[0] = (grad_mom[0] - vel[0] * drho) / r;
    grad_vel[1] = (grad_mom[1] - vel[1] * drho) / r;
    grad_vel[2] = (grad_mom[2] - vel[2] * drho) / r;

    RealTensor grad_vel_transpose = vtkm::MatrixTranspose(grad_vel);
    FloatDefault vel_div = grad_vel(0, 0) + grad_vel(1, 1) + grad_vel(2, 2);

    RealTensor tau;
    tau[0] = grad_vel[0] + grad_vel_transpose[0];
    tau[1] = grad_vel[1] + grad_vel_transpose[1];
    tau[2] = grad_vel[2] + grad_vel_transpose[2];

    tau(0, 0) -= 2. / 3 * vel_div;
    tau(1, 1) -= 2. / 3 * vel_div;
    tau(2, 2) -= 2. / 3 * vel_div;

    FloatDefault vis = 1;
    for (size_t i = 0; i < 3; i++)
    {
      for (size_t j = 0; j < 3; j++)
      {
        tau(i, j) *= vis / _reynolds;
      }
    }

    Vec3f grad_enthalpy =
      (duh[4] - uh[4] / uh[0] * duh[0]) / r - MatrixMultiply(vtkm::MatrixTranspose(grad_vel), vel);

    for (size_t i = 0; i < 3; i++)
    {
      grad_enthalpy[i] *= (vis / _reynolds) * (_gamma / _prandtl);
    }

    vis_term[0] = vtkm::TypeTraits<Vec3f>::ZeroInitialization();
    vis_term[1] = tau[0];
    vis_term[2] = tau[1];
    vis_term[3] = tau[2];
    vis_term[4] = MatrixMultiply(tau, vel) + grad_enthalpy;
  }
};

class EulerSystem : public SystemBase
{
public:
  static constexpr UInt8 NUM_VARIABLES = 5;
  using VType = Vec<FloatDefault, NUM_VARIABLES>;


  struct ComputeResidual : vtkm::worklet::WorkletCellNeighborhood
  {
    VTKM_CONT
    ComputeResidual(FloatDefault t)
      : _t(t)
    {
    }

    using ControlSignature =
      void(CellSetIn, FieldInNeighborhood, FieldOut Rhs, FieldIn centroid, FieldIn spacing);
    using ExecutionSignature = void(Boundary, _2, _3, _4, _5);

    template<typename NeighIn, typename RhsType>
    VTKM_EXEC void operator()(const vtkm::exec::BoundaryState& boundary,
                              const NeighIn& value,
                              RhsType& rhs,
                              const Vec3f& centroid,
                              const Vec3f& spacing) const
    {
      //rhs = vtkm::TypeTraits<RhsType>::ZeroInitialization();
      FloatDefault penalty = -8;
      Vec3f normal{ 0, 0, 0 };
      DataPack<RhsType::NUM_COMPONENTS> dl, dr, lift;
      {
        //auto& ul = value.Get(0, 0, 0);
        //auto ur =  value.Get(-1, 0, 0);
        dl.uh = value.Get(0, 0, 0);
        dr.uh = value.Get(-1, 0, 0);
        RhsType lflux, rflux;
        if (boundary.IJK[0] == 0)
        {
          Vec3f p = centroid;
          p[0] = p[0] - 0.5 * spacing[0];
          //ur = ul;
          //dr.uh = value.Get(0, 0, 0);
          dr.uh = valueExactCoutteFlow(_t, p);
        }
        dl.InvisTerm();
        dr.InvisTerm();

        normal = Vec3f{ -1, 0, 0 };
        for (size_t i = 0; i < RhsType::NUM_COMPONENTS; ++i)
        {
          lift.uh[i] = (dl.uh[i] + dr.uh[i]) / 2;
          lift.duh[i] = (dl.uh[i] - dr.uh[i]) / 2 * normal;
        }

        RiemannFlux(dl, dr, normal, lflux);
        lift.ViscousTerm();
        //std::cout << lift.vis_term[0] << std::endl;
        //std::cout << lift.vis_term[1] << std::endl;
        //std::cout << lift.vis_term[2] << std::endl;
        //std::cout << lift.vis_term[3] << std::endl;
        //std::cout << lift.vis_term[4] << std::endl;

        for (size_t i = 0; i < RhsType::NUM_COMPONENTS; i++)
        {
          lflux[i] += penalty * vtkm::Dot(lift.vis_term[i], normal) / spacing[0];
        }


        dl.uh = value.Get(0, 0, 0);
        dr.uh = value.Get(1, 0, 0);
        if (boundary.IJK[0] == boundary.PointDimensions[0] - 1)
        {
          Vec3f p = centroid;
          p[0] = p[0] + 0.5 * spacing[0];
          //ur = ul;
          //dr.uh = value.Get(0, 0, 0);
          dr.uh = valueExactCoutteFlow(_t, p);
        }
        dl.InvisTerm();
        dr.InvisTerm();


        normal = Vec3f{ 1, 0, 0 };
        for (size_t i = 0; i < RhsType::NUM_COMPONENTS; ++i)
        {
          lift.uh[i] = (dl.uh[i] + dr.uh[i]) / 2;
          lift.duh[i] = (dl.uh[i] - dr.uh[i]) / 2 * normal;
        }

        RiemannFlux(dl, dr, normal, rflux);
        lift.ViscousTerm();
        for (size_t i = 0; i < RhsType::NUM_COMPONENTS; i++)
        {
          rflux[i] += penalty * vtkm::Dot(lift.vis_term[i], normal) / spacing[0];
        }

        //RiemannFlux(ul, ur, Vec3f{ 1, 0, 0 }, rflux);
        rhs = (rflux + lflux) / spacing[0];
      }

      {
        //auto& ul = value.Get(0, 0, 0);
        //auto ur = value.Get(0, -1, 0);
        dl.uh = value.Get(0, 0, 0);
        dr.uh = value.Get(0, -1, 0);
        RhsType lflux, rflux;
        if (boundary.IJK[1] == 0)
        {
          Vec3f p = centroid;
          p[1] = p[1] - 0.5 * spacing[1];
          // dr.uh = value.Get(0, 0, 0);
          dr.uh = valueExactCoutteFlow(_t, p);
        }

        //dl.uh = ul;
        //dr.uh = ur;
        dl.InvisTerm();
        dr.InvisTerm();

        normal = Vec3f{ 0, -1, 0 };
        for (size_t i = 0; i < RhsType::NUM_COMPONENTS; ++i)
        {
          lift.uh[i] = (dl.uh[i] + dr.uh[i]) / 2;
          lift.duh[i] = (dl.uh[i] - dr.uh[i]) / 2 * normal;
        }

        RiemannFlux(dl, dr, normal, lflux);
        lift.ViscousTerm();
        for (size_t i = 0; i < RhsType::NUM_COMPONENTS; i++)
        {
          lflux[i] += penalty * vtkm::Dot(lift.vis_term[i], normal) / spacing[0];
        }

        //ur = value.Get(0, 1, 0);
        dl.uh = value.Get(0, 0, 0);
        dr.uh = value.Get(0, 1, 0);
        if (boundary.IJK[1] == boundary.PointDimensions[1] - 1)
        {
          Vec3f p = centroid;
          p[1] = p[1] + 0.5 * spacing[1];
          //dr.uh = value.Get(0, 0, 0);
          dr.uh = valueExactCoutteFlow(_t, p);
        }
        dl.InvisTerm();
        dr.InvisTerm();

        normal = Vec3f{ 0, 1, 0 };
        for (size_t i = 0; i < RhsType::NUM_COMPONENTS; ++i)
        {
          lift.uh[i] = (dl.uh[i] + dr.uh[i]) / 2;
          lift.duh[i] = (dl.uh[i] - dr.uh[i]) / 2 * normal;
        }

        RiemannFlux(dl, dr, normal, rflux);
        lift.ViscousTerm();
        for (size_t i = 0; i < RhsType::NUM_COMPONENTS; i++)
        {
          rflux[i] += penalty * vtkm::Dot(lift.vis_term[i], normal) / spacing[0];
        }

        rhs += (rflux + lflux) / spacing[1];
      }
    }

    VTKM_EXEC
    vtkm::Vec<FloatDefault, 5> valueExactIsoVortex(FloatDefault t, const Vec3f& p) const
    {
      Vec3f _vel{ 1, 1, 0 };
      FloatDefault x = p[0] - _vel[0] * t;
      FloatDefault y = p[1] - _vel[1] * t;
      FloatDefault z = p[2] - _vel[2] * t;
      FloatDefault gam = 1.4, gamm1 = gam - 1, epi = 5.;
      FloatDefault xb, yb, r2;

      FloatDefault PI = vtkm::Pi();
      xb = x + 5;
      yb = y + 5;
      r2 = xb * xb + yb * yb;

      FloatDefault u = _vel[0] + epi / (2. * PI) * exp(0.5 * (1. - r2)) * (-yb);
      FloatDefault v = _vel[1] + epi / (2. * PI) * exp(0.5 * (1. - r2)) * xb;
      FloatDefault T = 1. - gamm1 * epi * epi / (8 * gam * PI * PI) * exp(1. - r2);
      FloatDefault density = pow(T, 1 / gamm1);
      FloatDefault pre = pow(density, gam);

      vtkm::Vec<FloatDefault, 5> value;
      value[0] = density;
      value[1] = density * u;
      value[2] = density * v;
      value[3] = density * 0;
      value[4] = pre / gamm1 + 0.5 * density * (u * u + v * v);

      return value;
    }

    VTKM_EXEC
    vtkm::Vec<FloatDefault, 5> valueExactCoutteFlow(FloatDefault t, const Vec3f& p) const
    {
      FloatDefault _gamma = 1.4, _prandtl = 0.72, _mach = 0.2;

      FloatDefault x = p[0];
      FloatDefault y = p[1];
      FloatDefault z = p[2];

      FloatDefault T2 = 0.85, T1 = 0.8;
      FloatDefault tem = T1 + (T2 - T1) * y / 2 +
        0.5 * _prandtl * (_gamma - 1) * _mach * _mach * y / 2 * (1 - y / 2);
      FloatDefault pre = 1. / (_gamma * _mach * _mach);
      FloatDefault density = pre * _gamma * _mach * _mach / tem;
      FloatDefault u = y / 2.0, v = 0, w = 0;
      Vec3f mom{ density * u, density * v, density * w };
      FloatDefault rhoe = pre / (_gamma - 1) + 0.5 * density * (u * u + v * v + w * w);

      vtkm::Vec<FloatDefault, 5> value;
      value[0] = density;
      value[1] = density * u;
      value[2] = density * v;
      value[3] = density * w;
      value[4] = rhoe;

      return value;
    }
    template<typename DataPackType, typename VType>
    VTKM_EXEC void RiemannFlux(const DataPackType& dl,
                               const DataPackType& dr,
                               const Vec3f& normal,
                               VType& riemann_flux) const
    {
      for (IdComponent eq = 0; eq < VType::NUM_COMPONENTS; eq++)
      {
        riemann_flux[eq] =
          vtkm::Dot(dl.inv_term[eq] + dr.inv_term[eq], normal) / 2 + 0.6 * (dl.uh[eq] - dr.uh[eq]);
      }
    }
    FloatDefault _t;
  };

  struct Advanced : public vtkm::worklet::WorkletMapField
  {
    VTKM_CONT
    Advanced(FloatDefault dt)
      : _dt(dt)
    {
    }

    using ControlSignature = void(FieldInOut var, FieldIn rhs);
    using ExecutionSignature = void(_1, _2);

    template<typename T>
    VTKM_EXEC void operator()(T& u, const T& rhs) const
    {
      u += -_dt * rhs;
    }

    FloatDefault _dt;
  };


public:
  struct WriteVariable : public vtkm::worklet::WorkletMapField
  {
    VTKM_CONT
    WriteVariable(UInt8 eq)
      : _eq(eq)
    {
    }
    using ControlSignature = void(FieldIn solution, FieldOut field);
    using ExecutionSignature = void(_1, _2);

    template<typename SolType, typename FieldType>
    VTKM_EXEC void operator()(const SolType& sol, FieldType& field) const
    {
      field = sol[_eq];
    }

    UInt8 _eq;
  };

  EulerSystem(MeshBase& mesh)
    : SystemBase(mesh)
  {
  }

  virtual ~EulerSystem(){};

  void AddVariable() {}

  void Init(FloatDefault dt, Id steps)
  {
    this->_dt = dt;
    this->_total_steps = steps;
  }

  void InitCondition()
  {
    vtkm::cont::Timer timer;
    timer.Start();
    vtkm::cont::Invoker invoker;
    invoker(CoutteFlowIC{}, _mesh.GetCellSet(), _solution, _mesh.GetCoordinate());

    Success("初始条件设置完成", timer.GetElapsedTime(), "s");
  }

  void Solve()
  {
    vtkm::cont::Timer wall_timer;
    wall_timer.Start();
    for (_cur_step = 0; _cur_step <= _total_steps; _cur_step++)
    {
      //Output();

      vtkm::cont::Timer timer;
      fmt::print("Step{0:<4d} Physical Time = {1:.4f}  dt = {2:.3e}\n", _cur_step, _t, _dt);

      timer.Start();
      Update();
      fmt::print(
        fmt::fg(fmt::color::blue), "         ElapsedTime: {0:.3f}\n", timer.GetElapsedTime());
      //timer.Stop();
    }
    fmt::print(
      fmt::fg(fmt::color::green), "Solve Finished. Wall Time: {}\n", wall_timer.GetElapsedTime());
  }
  void Update()
  {
    vtkm::cont::Timer timer;
    //timer.Start();
    vtkm::cont::Invoker{}(
      ComputeResidual{ _t }, _mesh.GetCellSet(), _solution, _rhs, _mesh._centroid, _mesh._spacing);
    //fmt::print(" ElapsedTime: {0:.3f}\n", timer.GetElapsedTime());

    //timer.Start();
    vtkm::cont::Invoker{}(Advanced{ _dt }, _solution, _rhs);
    //fmt::print(" ElapsedTime: {0:.3f}\n", timer.GetElapsedTime());

    auto result =
      vtkm::cont::Algorithm::Reduce(_rhs, vtkm::TypeTraits<VType>::ZeroInitialization());
    std::cout << "Res: " << result[0] / _rhs.GetNumberOfValues() << std::endl;
    this->_t += _dt;
  }


  void Output()
  {
    //if (_cur_step % 2 == 1)
    {
      vtkm::cont::Timer timer;
      timer.Start();
      fmt::print(fmt::fg(fmt::color::honey_dew), "Output ...");

      char tmp[8];
      sprintf(tmp, "%06d", _cur_step);
      std::string file = "system" + std::string(tmp) + ".vtk";
      vtkm::io::VTKDataSetWriter write(file);

      vtkm::cont::DataSet out_vtk;
      out_vtk.AddCoordinateSystem(_mesh.GetCoordinate());
      out_vtk.SetCellSet(_mesh.GetCellSet());
      for (UInt8 eq = 0; eq < NUM_VARIABLES; eq++)
      {
        out_vtk.AddCellField("variable" + std::to_string(eq), ArrayHandle<FloatDefault>{});
      }
      for (UInt8 eq = 0; eq < NUM_VARIABLES; eq++)
      {
        Invoker{}(WriteVariable{ eq }, _solution, out_vtk.GetField(eq));
      }
      write.WriteDataSet(out_vtk);

      Success("输出完成", timer.GetElapsedTime(), "s");
    }
  }


  const MeshBase& GetMesh() const { return _mesh; }
  const ArrayHandle<VType>& Solution() const { return _solution; }

private:
  ArrayHandle<VType> _solution;
  ArrayHandle<VType> _rhs;
  Id _cur_step = 0, _total_steps = 30;
  FloatDefault _t = 0, _dt = 0.0001;
};
